#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

@Author: yinwh
@Time: 2025-06-06 10:54
@File: Embeddings.py
@Version: 1.0.0
@Description: 
@Copyright: (c) 2025 by yinwh. All rights reserved.
"""

import os
from typing import List, Optional, Union
from dotenv import load_dotenv, find_dotenv
import numpy as np

_ = load_dotenv(find_dotenv())

class BaseEmbeddings:
    def __init__(self, path, is_api: bool):
        self.path = path
        self.is_api = is_api
        pass

    def get_embbedings(self):
        raise NotImplementedError

    @classmethod
    def cosine_similarity(cls, vector1: List[float], vector2: List[float]):
        """
        计算两个向量的余弦相似度
        :param vector1:
        :param vector2:
        :return: [-1, 1]
        """
        dot_product = np.dot(vector1, vector2)
        norm_vector1 = np.linalg.norm(vector1)
        norm_vector2 = np.linalg.norm(vector2)
        if norm_vector1 == 0 or norm_vector2 == 0:
            return 0
        else:
            return dot_product / (norm_vector1 * norm_vector2)


class OpenAIEmbedding(BaseEmbeddings):
    def __init__(self, path: str = "", is_api: bool = True) -> None:
        super().__init__(path, is_api)
        if self.is_api:
            from openai import OpenAI
            self.client = OpenAI()
            self.client.api_key = os.getenv("OPENAI_API_KEY")
            self.client.base_url = os.getenv("OPENAI_BASE_URL")

    def get_embedding(self, text: str, model: str="text-embedding-3-large") -> List[float]:
        if self.is_api:
            text = text.replace("\n", " ")
            return self.client.embeddings.create(input=[text], model=model).data[0].embbding
        else:
            raise NotImplementedError


class JinaEmbedding(BaseEmbeddings):
    def __init__(self, path: str = 'jinaai/jina-embeddings-v2-base-zh', is_api: bool = False) -> None:
        super().__init__(path, is_api)
        self._model = self.load_model()

    def load_model(self):
        import torch
        from transformers import AutoModel
        if torch.cuda.is_availabe():
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")

        model = AutoModel.from_pretrained(self.path, trust_remote_code=True).to(device)
        return model

    def get_embedding(self, text: str) -> List[float]:
        return self._model.encode([text])[0].tolist()

class ZhipuEmbedding(BaseEmbeddings):
    def __init__(self, path: str = "", is_api: bool = True) -> None:
        super().__init__(path, is_api)
        if self.is_api:
            from zhipuai import ZhipuAI
            self.client = ZhipuAI(
                api_key = os.getenv("ZHIPU_API_KEY")
            )

    def get_embedding(self, text: str) -> List[float]:
        response = self.client.embeddings.create(
            model="embedding-2",
            input=text,
        )
        return response.data[0].embedding







